# -*- coding: utf-8 -*-

# import tushare as ts, os, pandas as pd
# ts.set_token('9a7b38fcb04482d22b21c79728fb21ea1a05fde3d9ac779da086801d')          # 官网免费注册
# pro = ts.pro_api()

# stocks = pro.stock_basic(exchange='', list_status='L')['ts_code']
# for code in stocks[:5]:
#     try:
#         df = pro.daily_basic(ts_code=code, start_date='20220101', end_date='20231231')
#         df.to_csv(f'./cn_annual_2022_2022/{code}.csv', index=False)
#     except Exception as e:
#         continue

import os, time, requests, pandas as pd, tushare as ts
from tqdm import tqdm

# ---------- 0. 登录 ----------
ts.set_token('9a7b38fcb04482d22b21c79728fb21ea1a05fde3d9ac779da086801d')
pro = ts.pro_api()

# ---------- 1. A 股列表（剔除退市） ----------
stocks = pro.stock_basic(exchange='', list_status='L', fields='ts_code,symbol,name')
tickers = stocks['ts_code'].tolist()

# ---------- 2. 输出目录 ----------
OUT_DIR = './cn_annual_2003_2022'
os.makedirs(OUT_DIR, exist_ok=True)

# ---------- 3. 抓取 & 下载 ----------
headers = {'User-Agent': 'Mozilla/5.0'}
records = []

for ts_code in tqdm(tickers, desc='Tickers'):
    # 一次取 20 年，防止接口分页
    df_ann = pro.anns(ts_code=ts_code,
                      start_date='20220101',
                      end_date='20231231',
                      fields='ts_code,ann_date,title,f_path')
    # 只保留“年度报告”、“年报”、“年度报告（更新后）”等
    mask = df_ann['title'].str.contains('年度报告')
    df_ann = df_ann[mask]

    for _, row in df_ann.iterrows()[:5]:
        url = row['f_path']
        year  = row['ann_date'][:4]  # 2022-03-31 -> 2022
        fname = f"{ts_code}_{year}.pdf"
        fpath = os.path.join(OUT_DIR, fname)

        if os.path.exists(fpath):       # 断点续下
            continue

        try:
            r = requests.get(url, headers=headers, timeout=30)
            r.raise_for_status()
            with open(fpath, 'wb') as f:
                f.write(r.content)
            records.append({
                'ts_code': ts_code,
                'year': int(year),
                'title': row['title'],
                'path': os.path.abspath(fpath),
                'ann_date': row['ann_date']
            })
        except Exception as e:
            print('❌', ts_code, year, e)

        time.sleep(0.2)  # 限速

# ---------- 4. 生成索引 ----------
pd.DataFrame(records).to_csv('cn_annual_index.csv', index=False)
print('✅ 完成！共下载', len(records), '份年报')
